from unicodedata import numeric
import cv2
import numpy
import torch
import pickle 
import config as cfg
from tqdm import tqdm
from os import listdir, name
from os.path import join
from numpy import array, zeros, save
from models.facenet.facenet import Facenet
from face_recognition import face_locations, load_image_file
from sklearn.neighbors import KNeighborsClassifier

def encoding(img_path):
    paths = listdir(img_path)
    features = []
    names = []
    warn_msgs = []

    # 加载模型
    path = 'models/facenet/facenet_mobilenet.pth'
    model = Facenet(mode='predict').eval()
    state_dict = torch.load(path, map_location='cpu')
    model.load_state_dict(state_dict, strict=False)

    # 遍历人脸图片
    for i, path in zip(tqdm(range(len(paths))), paths):
        image = load_image_file(join(img_path, path))
        small_img = cv2.resize(image, (0, 0), fx=0.25, fy=0.25)
        faces = []
        boxes = face_locations(small_img, model='hog')
        if len(boxes) != 1:
            # 人脸不等于1， 存入消息表，运行完成后统一输出。
            warn_msgs.append("file {0} can't found face, or found more than 1 face".format(join(img_path, path)))
            continue
        else:
            # 将人脸特征向量与人名写入列表。
            box = boxes[0]
            face = image[box[0]:box[2], box[3]:box[1]]
            face = cv2.cvtColor(face, cv2.COLOR_RGB2BGR)
            faces.append(cv2.resize(face, (160, 160)))
            faces = array(faces).transpose(0, 3, 1, 2)
            faces = torch.from_numpy(faces).type(torch.FloatTensor)
            feature = model(faces).detach().numpy()
            features.append(feature[0])
            name = path.split('_')[0]
            names.append(name)
        
    names = array(names)

    return features, names, warn_msgs


if __name__ == '__main__':
    print('start encoding faces')
    features, names, warn_msgs = encoding(cfg.pictures_of_konw)
    neigh = KNeighborsClassifier(n_neighbors=3, algorithm='ball_tree', weights='distance')
    neigh.fit(features, names)
    with open('data/knn.pickle', 'wb') as f:
        pickle.dump(neigh, f)
    numpy.save('data/names.npy', names)
    numpy.save('data/features.npy', features)
    print('encode faces complete, take {} warnings:'.format(len(warn_msgs)))
    for msg in warn_msgs:
        print('\twarnings:{}'.format(msg))